// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/functions/lib/window/tests/WindowTestBase.h>
#include <pollux/functions/sparksql/window/window_functions_registration.h>

using namespace kumo::pollux::exec::test;

namespace kumo::pollux::window::test {

namespace {

static const std::vector<std::string> kSparkWindowFunctions = {
    std::string("nth_value(c0, 1)"),
    std::string("nth_value(c0, c3)"),
    std::string("row_number()"),
    std::string("rank()"),
    std::string("dense_rank()"),
    std::string("ntile(c3)"),
    std::string("ntile(1)"),
    std::string("ntile(7)"),
    std::string("ntile(10)"),
    std::string("ntile(16)")};

struct SparkWindowTestParam {
  const std::string function;
  const std::string overClause;
};

class SparkWindowTestBase : public WindowTestBase {
 protected:
  explicit SparkWindowTestBase(const SparkWindowTestParam& testParam)
      : function_(testParam.function), overClause_(testParam.overClause) {}

  void testWindowFunction(const std::vector<RowVectorPtr>& vectors) {
    WindowTestBase::testWindowFunction(vectors, function_, {overClause_});
  }

  void SetUp() override {
    WindowTestBase::SetUp();
    WindowTestBase::options_.parseIntegerAsBigint = false;
    pollux::functions::window::sparksql::registerWindowFunctions("");
  }

  const std::string function_;
  const std::string overClause_;
};

std::vector<SparkWindowTestParam> getSparkWindowTestParams() {
  std::vector<SparkWindowTestParam> params;

  for (auto function : kSparkWindowFunctions) {
    for (auto overClause : kOverClauses) {
      params.push_back({function, overClause});
    }
  }

  return params;
}

class SparkWindowTest
    : public SparkWindowTestBase,
      public testing::WithParamInterface<SparkWindowTestParam> {
 public:
  SparkWindowTest() : SparkWindowTestBase(GetParam()) {}
};

// Tests all functions with a dataset with uniform distribution of partitions.
TEST_P(SparkWindowTest, basic) {
  testWindowFunction({makeSimpleVector(40)});
}

// Tests all functions with a dataset with all rows in a single partition,
// but in 2 input vectors.
TEST_P(SparkWindowTest, singlePartition) {
  testWindowFunction(
      {makeSinglePartitionVector(40), makeSinglePartitionVector(50)});
}

// Tests all functions with a dataset in which all partitions have a single row.
TEST_P(SparkWindowTest, singleRowPartitions) {
  testWindowFunction({makeSingleRowPartitionsVector(40)});
}

// Tests all functions with a dataset with randomly generated data.
TEST_P(SparkWindowTest, randomInput) {
  testWindowFunction({makeRandomInputVector(30)});
}

// Run above tests for all combinations of rank function and over clauses.
POLLUX_INSTANTIATE_TEST_SUITE_P(
    SparkWindowTestInstantiation,
    SparkWindowTest,
    testing::ValuesIn(getSparkWindowTestParams()));

} // namespace
} // namespace kumo::pollux::window::test
